{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# Building Training Models\n",
    "\n",
    "[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/brainpy/brainpy/blob/master/docs_version2/tutorial_training/build_training_models.ipynb)\n",
    "[![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/brainpy/brainpy/blob/master/docs_version2/tutorial_training/build_training_models.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "In this section, we are going to talk about how to build models for training."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:26.480089Z",
     "start_time": "2025-10-06T05:19:22.091973Z"
    }
   },
   "source": [
    "import brainpy as bp\n",
    "import brainpy.math as bm\n",
    "\n",
    "bm.set_platform('cpu')\n",
    "\n",
    "bp.__version__"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'3.0.0'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 1
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Use built-in models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "``brainpy.DynamicalSystem`` provided in BrainPy can be used for model training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### ``mode`` settings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Some built-in models have implemented the training interface for their training. Users can instantiate these models by providing the parameter ``mode=brainpy.modes.training`` for training model customization."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "For example, ``brainpy.neurons.LIF`` is a model commonly used in computational simulation, but it can also be used in training."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:26.561638Z",
     "start_time": "2025-10-06T05:19:26.488608Z"
    }
   },
   "source": [
    "# Instantiate a LIF model for simulation\n",
    "\n",
    "lif = bp.neurons.LIF(1)\n",
    "lif.mode"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "NonBatchingMode"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 2
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:26.585550Z",
     "start_time": "2025-10-06T05:19:26.567233Z"
    }
   },
   "source": [
    "# Instantiate a LIF model for training.\n",
    "# In this mode, the model implement variables and functions\n",
    "# compatible with BrainPy's training interface.\n",
    "\n",
    "lif = bp.neurons.LIF(1, mode=bm.TrainingMode())\n",
    "lif.mode"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TrainingMode(batch_size=1)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 3
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "But some build-in models does not support training."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:26.711756Z",
     "start_time": "2025-10-06T05:19:26.631270Z"
    }
   },
   "source": [
    "try:\n",
    "    bp.layers.NVAR(1, 1, mode=bm.TrainingMode())\n",
    "except Exception as e:\n",
    "    print(type(e), e)"
   ],
   "outputs": [],
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "The ``mode`` can be used to control the weight types. Let's take a synaptic model for another example. For a non-trainable dense layer, the *weights* and *bias* are Array instances.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:26.931175Z",
     "start_time": "2025-10-06T05:19:26.717996Z"
    }
   },
   "source": [
    "l = bp.layers.Dense(3, 4, mode=bm.batching_mode)\n",
    "\n",
    "l.W"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(value=Array([[ 0.00514068, -0.38861045, -0.18789645, -0.20446473],\n",
       "                   [ 0.8905003 , -1.1153231 , -0.30391234, -1.1171588 ],\n",
       "                   [-0.41736808, -0.45995244,  0.513535  ,  0.64199084]],      dtype=float32),\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 5
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:27.131135Z",
     "start_time": "2025-10-06T05:19:26.949273Z"
    }
   },
   "source": [
    "l = bp.layers.Dense(3, 4, mode=bm.training_mode)\n",
    "\n",
    "l.W"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TrainVar(\n",
       "  value=ShapedArray(float32[3,4]),\n",
       "  _batch_axis=None,\n",
       "  axis_names=None\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 6
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Moreover, for some recurrent models, e.g., ``LSTM`` or ``GRU``, the ``state`` can be set to be trainable or not trainable by ``train_state`` argument. When setting ``train_state=True`` for the recurrent instance, a new attribute *.state2train* will be created."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:27.732610Z",
     "start_time": "2025-10-06T05:19:27.149998Z"
    }
   },
   "source": [
    "rnn = bp.dyn.RNNCell(1, 3, train_state=True, mode=bm.training_mode)\n",
    "\n",
    "rnn.state2train"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TrainVar(\n",
       "  value=ShapedArray(float32[3]),\n",
       "  _batch_axis=None,\n",
       "  axis_names=None\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 7
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Note the difference between the *.state2train* and the original *.state*:\n",
    "\n",
    "1. *.state2train* has no batch axis.\n",
    "2. When using `node.reset()` function, all values in the *.state* will be filled with *.state2train*."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.294808Z",
     "start_time": "2025-10-06T05:19:27.750466Z"
    }
   },
   "source": [
    "rnn.reset(batch_size=5)\n",
    "rnn.state"
   ],
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "Do not support type <class 'NoneType'>: None",
     "output_type": "error",
     "traceback": [
      "\u001B[31m---------------------------------------------------------------------------\u001B[39m",
      "\u001B[31mValueError\u001B[39m                                Traceback (most recent call last)",
      "\u001B[36mCell\u001B[39m\u001B[36m \u001B[39m\u001B[32mIn[8]\u001B[39m\u001B[32m, line 1\u001B[39m\n\u001B[32m----> \u001B[39m\u001B[32m1\u001B[39m rnn.reset(batch_size=\u001B[32m5\u001B[39m)\n\u001B[32m      2\u001B[39m rnn.state\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\BrainPy\\brainpy\\version2\\dynsys.py:178\u001B[39m, in \u001B[36mDynamicalSystem.reset\u001B[39m\u001B[34m(self, *args, **kwargs)\u001B[39m\n\u001B[32m    176\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m reset_state \u001B[38;5;129;01mis\u001B[39;00m \u001B[38;5;28;01mNone\u001B[39;00m:\n\u001B[32m    177\u001B[39m     \u001B[38;5;28;01mfrom\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34;01mbrainpy\u001B[39;00m\u001B[34;01m.\u001B[39;00m\u001B[34;01mversion2\u001B[39;00m\u001B[34;01m.\u001B[39;00m\u001B[34;01mhelpers\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[38;5;28;01mimport\u001B[39;00m reset_state\n\u001B[32m--> \u001B[39m\u001B[32m178\u001B[39m reset_state(\u001B[38;5;28mself\u001B[39m, *args, **kwargs)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\BrainPy\\brainpy\\version2\\helpers.py:61\u001B[39m, in \u001B[36mreset_state\u001B[39m\u001B[34m(target, *args, **kwargs)\u001B[39m\n\u001B[32m     59\u001B[39m \u001B[38;5;28;01mfor\u001B[39;00m node \u001B[38;5;129;01min\u001B[39;00m nodes:\n\u001B[32m     60\u001B[39m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;28mhasattr\u001B[39m(node.reset_state, \u001B[33m'\u001B[39m\u001B[33mreset_level\u001B[39m\u001B[33m'\u001B[39m):\n\u001B[32m---> \u001B[39m\u001B[32m61\u001B[39m         node.reset_state(*args, **kwargs)\n\u001B[32m     62\u001B[39m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m     63\u001B[39m         nodes_with_level.append(node)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\BrainPy\\brainpy\\version2\\dyn\\rates\\rnncells.py:114\u001B[39m, in \u001B[36mRNNCell.reset_state\u001B[39m\u001B[34m(self, batch_or_mode, **kwargs)\u001B[39m\n\u001B[32m    113\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34mreset_state\u001B[39m(\u001B[38;5;28mself\u001B[39m, batch_or_mode=\u001B[38;5;28;01mNone\u001B[39;00m, **kwargs):\n\u001B[32m--> \u001B[39m\u001B[32m114\u001B[39m     \u001B[38;5;28mself\u001B[39m.state.value = parameter(\u001B[38;5;28mself\u001B[39m._state_initializer, (batch_or_mode, \u001B[38;5;28mself\u001B[39m.num_out,), allow_none=\u001B[38;5;28;01mFalse\u001B[39;00m)\n\u001B[32m    115\u001B[39m     \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mself\u001B[39m.train_state:\n\u001B[32m    116\u001B[39m         \u001B[38;5;28mself\u001B[39m.state2train.value = parameter(\u001B[38;5;28mself\u001B[39m._state_initializer, \u001B[38;5;28mself\u001B[39m.num_out, allow_none=\u001B[38;5;28;01mFalse\u001B[39;00m)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\BrainPy\\brainpy\\version2\\initialize\\generic.py:83\u001B[39m, in \u001B[36mparameter\u001B[39m\u001B[34m(param, sizes, allow_none, allow_scalar, sharding)\u001B[39m\n\u001B[32m     80\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m param\n\u001B[32m     82\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m \u001B[38;5;28mcallable\u001B[39m(param):\n\u001B[32m---> \u001B[39m\u001B[32m83\u001B[39m     v = bm.jit(param,\n\u001B[32m     84\u001B[39m                static_argnums=\u001B[32m0\u001B[39m,\n\u001B[32m     85\u001B[39m                out_shardings=bm.sharding.get_sharding(sharding))(sizes)\n\u001B[32m     86\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m _check_var(v)  \u001B[38;5;66;03m# TODO: checking the Variable need to be traced\u001B[39;00m\n\u001B[32m     88\u001B[39m \u001B[38;5;28;01melif\u001B[39;00m \u001B[38;5;28misinstance\u001B[39m(param, (np.ndarray, jnp.ndarray)):\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_jit.py:97\u001B[39m, in \u001B[36m_get_jitted_fun.<locals>.jitted_fun\u001B[39m\u001B[34m(*args, **params)\u001B[39m\n\u001B[32m     94\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m fun.fun(*args, **params)\n\u001B[32m     96\u001B[39m \u001B[38;5;66;03m# compile the function and get the state trace\u001B[39;00m\n\u001B[32m---> \u001B[39m\u001B[32m97\u001B[39m state_trace = fun.get_state_trace_by_call(*args, **params)\n\u001B[32m     98\u001B[39m read_state_vals = state_trace.get_read_state_values(\u001B[38;5;28;01mTrue\u001B[39;00m)\n\u001B[32m    100\u001B[39m \u001B[38;5;66;03m# call the jitted function\u001B[39;00m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:689\u001B[39m, in \u001B[36mStatefulFunction.get_state_trace_by_call\u001B[39m\u001B[34m(self, compile_if_miss, *args, **kwargs)\u001B[39m\n\u001B[32m    671\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34mget_state_trace_by_call\u001B[39m(\u001B[38;5;28mself\u001B[39m, *args, compile_if_miss: \u001B[38;5;28mbool\u001B[39m = \u001B[38;5;28;01mTrue\u001B[39;00m, **kwargs) -> StateTraceStack:\n\u001B[32m    672\u001B[39m \u001B[38;5;250m    \u001B[39m\u001B[33;03m\"\"\"\u001B[39;00m\n\u001B[32m    673\u001B[39m \u001B[33;03m    Read the state trace of the function.\u001B[39;00m\n\u001B[32m    674\u001B[39m \n\u001B[32m   (...)\u001B[39m\u001B[32m    687\u001B[39m \u001B[33;03m        The state trace of the function.\u001B[39;00m\n\u001B[32m    688\u001B[39m \u001B[33;03m    \"\"\"\u001B[39;00m\n\u001B[32m--> \u001B[39m\u001B[32m689\u001B[39m     cache_key = \u001B[38;5;28mself\u001B[39m.get_arg_cache_key(*args, **kwargs, compile_if_miss=compile_if_miss)\n\u001B[32m    690\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m.get_state_trace(cache_key)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:839\u001B[39m, in \u001B[36mStatefulFunction.get_arg_cache_key\u001B[39m\u001B[34m(self, compile_if_miss, *args, **kwargs)\u001B[39m\n\u001B[32m    831\u001B[39m cache_key = hashabledict(\n\u001B[32m    832\u001B[39m     static_args=static_args,\n\u001B[32m    833\u001B[39m     dyn_args=dyn_args,\n\u001B[32m    834\u001B[39m     static_kwargs=static_kwargs,\n\u001B[32m    835\u001B[39m     dyn_kwargs=dyn_kwargs,\n\u001B[32m    836\u001B[39m )\n\u001B[32m    838\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m cache_key \u001B[38;5;129;01mnot\u001B[39;00m \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mself\u001B[39m._cached_state_trace \u001B[38;5;129;01mand\u001B[39;00m compile_if_miss:\n\u001B[32m--> \u001B[39m\u001B[32m839\u001B[39m     \u001B[38;5;28mself\u001B[39m.make_jaxpr(*args, **kwargs)\n\u001B[32m    841\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m cache_key\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:950\u001B[39m, in \u001B[36mStatefulFunction.make_jaxpr\u001B[39m\u001B[34m(self, *args, **kwargs)\u001B[39m\n\u001B[32m    948\u001B[39m         \u001B[38;5;28mself\u001B[39m._cached_out_shapes.pop(cache_key, \u001B[38;5;28;01mNone\u001B[39;00m)\n\u001B[32m    949\u001B[39m         \u001B[38;5;28mself\u001B[39m._cached_jaxpr.pop(cache_key, \u001B[38;5;28;01mNone\u001B[39;00m)\n\u001B[32m--> \u001B[39m\u001B[32m950\u001B[39m         \u001B[38;5;28;01mraise\u001B[39;00m e\n\u001B[32m    952\u001B[39m \u001B[38;5;28;01mreturn\u001B[39;00m \u001B[38;5;28mself\u001B[39m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:928\u001B[39m, in \u001B[36mStatefulFunction.make_jaxpr\u001B[39m\u001B[34m(self, *args, **kwargs)\u001B[39m\n\u001B[32m    926\u001B[39m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m    927\u001B[39m         dyn_kwargs[k] = v\n\u001B[32m--> \u001B[39m\u001B[32m928\u001B[39m jaxpr, (out_shapes, state_shapes) = _make_jaxpr(\n\u001B[32m    929\u001B[39m     functools.partial(\n\u001B[32m    930\u001B[39m         \u001B[38;5;28mself\u001B[39m._wrapped_fun_to_eval,\n\u001B[32m    931\u001B[39m         cache_key,\n\u001B[32m    932\u001B[39m         static_kwargs,\n\u001B[32m    933\u001B[39m         return_only_write=\u001B[38;5;28mself\u001B[39m.return_only_write,\n\u001B[32m    934\u001B[39m     ),\n\u001B[32m    935\u001B[39m     static_argnums=\u001B[38;5;28mself\u001B[39m.static_argnums,\n\u001B[32m    936\u001B[39m     axis_env=\u001B[38;5;28mself\u001B[39m.axis_env,\n\u001B[32m    937\u001B[39m     return_shape=\u001B[38;5;28;01mTrue\u001B[39;00m,\n\u001B[32m    938\u001B[39m     abstracted_axes=\u001B[38;5;28mself\u001B[39m.abstracted_axes,\n\u001B[32m    939\u001B[39m )(*args, **dyn_kwargs)\n\u001B[32m    940\u001B[39m \u001B[38;5;66;03m# returns\u001B[39;00m\n\u001B[32m    941\u001B[39m \u001B[38;5;28mself\u001B[39m._cached_jaxpr_out_tree.set(cache_key, jax.tree.structure((out_shapes, state_shapes)))\n",
      "    \u001B[31m[... skipping hidden 1 frame]\u001B[39m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:1399\u001B[39m, in \u001B[36m_make_jaxpr.<locals>.make_jaxpr_f\u001B[39m\u001B[34m(*args, **kwargs)\u001B[39m\n\u001B[32m   1397\u001B[39m         jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info_)\n\u001B[32m   1398\u001B[39m     \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m-> \u001B[39m\u001B[32m1399\u001B[39m         jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f)\n\u001B[32m   1400\u001B[39m closed_jaxpr = ClosedJaxpr(jaxpr, consts)\n\u001B[32m   1401\u001B[39m \u001B[38;5;28;01mif\u001B[39;00m return_shape:\n",
      "    \u001B[31m[... skipping hidden 6 frame]\u001B[39m\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\brainstate\\brainstate\\transform\\_make_jaxpr.py:874\u001B[39m, in \u001B[36mStatefulFunction._wrapped_fun_to_eval\u001B[39m\u001B[34m(self, cache_key, static_kwargs, return_only_write, *args, **dyn_kwargs)\u001B[39m\n\u001B[32m    872\u001B[39m \u001B[38;5;28mself\u001B[39m._cached_state_trace.set(cache_key, state_trace)\n\u001B[32m    873\u001B[39m \u001B[38;5;28;01mwith\u001B[39;00m state_trace:\n\u001B[32m--> \u001B[39m\u001B[32m874\u001B[39m     out = \u001B[38;5;28mself\u001B[39m.fun(*args, **dyn_kwargs, **static_kwargs)\n\u001B[32m    875\u001B[39m     state_values = (\n\u001B[32m    876\u001B[39m         state_trace.get_write_state_values(\u001B[38;5;28;01mTrue\u001B[39;00m)\n\u001B[32m    877\u001B[39m         \u001B[38;5;28;01mif\u001B[39;00m return_only_write \u001B[38;5;28;01melse\u001B[39;00m\n\u001B[32m    878\u001B[39m         state_trace.get_state_values()\n\u001B[32m    879\u001B[39m     )\n\u001B[32m    880\u001B[39m state_trace.recovery_original_values()\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\BrainPy\\brainpy\\version2\\math\\object_transform\\_utils.py:68\u001B[39m, in \u001B[36mwarp_to_no_state_input_output.<locals>.wrapper\u001B[39m\u001B[34m(*args, **kwargs)\u001B[39m\n\u001B[32m     65\u001B[39m \u001B[38;5;129m@wraps\u001B[39m(fn)\n\u001B[32m     66\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34mwrapper\u001B[39m(*args, **kwargs):\n\u001B[32m     67\u001B[39m     args, kwargs = jax.tree.map(_remove_state, (args, kwargs), is_leaf=\u001B[38;5;28;01mlambda\u001B[39;00m x: \u001B[38;5;28misinstance\u001B[39m(x, brainstate.State))\n\u001B[32m---> \u001B[39m\u001B[32m68\u001B[39m     out = fn(*args, **kwargs)\n\u001B[32m     69\u001B[39m     out = jax.tree.map(_remove_state, out, is_leaf=\u001B[38;5;28;01mlambda\u001B[39;00m x: \u001B[38;5;28misinstance\u001B[39m(x, brainstate.State))\n\u001B[32m     70\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m out\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\BrainPy\\brainpy\\version2\\initialize\\regular_inits.py:21\u001B[39m, in \u001B[36mZeroInit.__call__\u001B[39m\u001B[34m(self, shape, dtype)\u001B[39m\n\u001B[32m     20\u001B[39m \u001B[38;5;28;01mdef\u001B[39;00m\u001B[38;5;250m \u001B[39m\u001B[34m__call__\u001B[39m(\u001B[38;5;28mself\u001B[39m, shape, dtype=\u001B[38;5;28;01mNone\u001B[39;00m):\n\u001B[32m---> \u001B[39m\u001B[32m21\u001B[39m     shape = [tools.size2num(d) \u001B[38;5;28;01mfor\u001B[39;00m d \u001B[38;5;129;01min\u001B[39;00m shape]\n\u001B[32m     22\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m bm.zeros(shape, dtype=dtype)\n",
      "\u001B[36mFile \u001B[39m\u001B[32mD:\\codes\\projects\\BrainPy\\brainpy\\version2\\tools\\others.py:82\u001B[39m, in \u001B[36msize2num\u001B[39m\u001B[34m(size)\u001B[39m\n\u001B[32m     80\u001B[39m     \u001B[38;5;28;01mreturn\u001B[39;00m a\n\u001B[32m     81\u001B[39m \u001B[38;5;28;01melse\u001B[39;00m:\n\u001B[32m---> \u001B[39m\u001B[32m82\u001B[39m     \u001B[38;5;28;01mraise\u001B[39;00m \u001B[38;5;167;01mValueError\u001B[39;00m(\u001B[33mf\u001B[39m\u001B[33m'\u001B[39m\u001B[33mDo not support type \u001B[39m\u001B[38;5;132;01m{\u001B[39;00m\u001B[38;5;28mtype\u001B[39m(size)\u001B[38;5;132;01m}\u001B[39;00m\u001B[33m: \u001B[39m\u001B[38;5;132;01m{\u001B[39;00msize\u001B[38;5;132;01m}\u001B[39;00m\u001B[33m'\u001B[39m)\n",
      "\u001B[31mValueError\u001B[39m: Do not support type <class 'NoneType'>: None"
     ]
    }
   ],
   "execution_count": 8
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Naming a node"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "For convenience, you can name a layer by specifying the name keyword argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.299173400Z",
     "start_time": "2023-04-15T17:18:07.780260Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dense(name=hidden_layer, num_in=128, num_out=100, mode=NonBatchingMode)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bp.layers.Dense(128, 100, name='hidden_layer')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Initializing parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Many models have their parameters. We can set the parameter of a model with the following methods."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "- **Arrays**\n",
    "\n",
    "If an array is provided, this is used unchanged as the parameter variable. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.300765500Z",
     "start_time": "2023-04-15T17:18:07.943925Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10, 50)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l = bp.layers.Dense(10, 50, W_initializer=bm.random.normal(0, 0.01, size=(10, 50)))\n",
    "\n",
    "l.W.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "- **Callable function**\n",
    "\n",
    "If a callable function (which receives a ``shape`` argument) is provided, the callable will be called with the desired shape to generate suitable initial parameter values. The variable is then initialized with those values. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.300765500Z",
     "start_time": "2023-04-15T17:18:08.132222Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(20, 30)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def init(shape):\n",
    "    return bm.random.random(shape)\n",
    "\n",
    "l = bp.layers.Dense(20, 30, W_initializer=init)\n",
    "\n",
    "l.W.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "- **Instance of** ``brainpy.init.Initializer``\n",
    "\n",
    "If a ``brainpy.init.Initializer`` instance is provided, the initial parameter values will be generated with the desired shape by using the Initializer instance. For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.300765500Z",
     "start_time": "2023-04-15T17:18:08.273000Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(20, 30)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "l = bp.layers.Dense(20, 30, W_initializer=bp.init.Normal(0.01))\n",
    "\n",
    "l.W.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "The weight matrix $W$ of this dense layer will be initialized using samples from a normal distribution with standard deviation 0.01 (see [brainpy.init](../apis/auto/initialize.rst) for more information)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "- **None parameter**\n",
    "\n",
    "Some types of parameter variables can also be set to ``None`` at initialization (e.g. biases). In that case, the parameter variable will be omitted. For example, creating a dense layer without biases is done as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.300765500Z",
     "start_time": "2023-04-15T17:18:08.429894Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "None\n"
     ]
    }
   ],
   "source": [
    "l = bp.layers.Dense(20, 100, b_initializer=None)\n",
    "\n",
    "print(l.b)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Customize your models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Customizing your training models is simple. You just need to subclass ``brainpy.DynamicalSystem``, and implement its ``update()`` and ``reset_state()`` functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Here, we demonstrate the model customization using two examples. The first is a recurrent layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.300765500Z",
     "start_time": "2023-04-15T17:18:08.594351Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class RecurrentLayer(bp.DynamicalSystemNS):\n",
    "    def __init__(self, num_in, num_out):\n",
    "        super(RecurrentLayer, self).__init__()\n",
    "\n",
    "        bp.check.is_subclass(self.mode, (bm.TrainingMode, bm.BatchingMode))\n",
    "\n",
    "        # define parameters\n",
    "        self.num_in = num_in\n",
    "        self.num_out = num_out\n",
    "\n",
    "        # define variables\n",
    "        self.state = bm.Variable(bm.zeros(1, num_out), batch_axis=0)\n",
    "\n",
    "        # define weights\n",
    "        self.win = bm.TrainVar(bm.random.normal(0., 1./num_in ** 0.5, size=(num_in, num_out)))\n",
    "        self.wrec = bm.TrainVar(bm.random.normal(0., 1./num_out ** 0.5, size=(num_out, num_out)))\n",
    "\n",
    "    def reset_state(self, batch_size):\n",
    "        # this function defines how to reset the mode states\n",
    "        self.state.value = bm.zeros((batch_size, self.num_out))\n",
    "\n",
    "    def update(self, x):\n",
    "        # this function defined how the model update its state and produce its output\n",
    "        out = bm.dot(x, self.win) + bm.dot(self.state, self.wrec)\n",
    "        self.state.value = bm.tanh(out)\n",
    "        return self.state.value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "This simple example illustrates many features essential for a training model. ``reset_state()`` function defines how to reset model states, which will be called at the first time step; ``update()`` function defines how the model states are evolving, which will be called at every time step."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Another example is the dropout layer, which can be useful to demonstrate how to define a model with multiple behaviours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.300765500Z",
     "start_time": "2023-04-15T17:18:08.610800Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class Dropout(bp.DynamicalSystemNS):\n",
    "  def __init__(self, prob: float, seed: int = None, name: str = None):\n",
    "    super(Dropout, self).__init__(name=name)\n",
    "\n",
    "    bp.check.is_subclass(self.mode, (bm.TrainingMode, bm.BatchingMode, bm.NonBatchingMode))\n",
    "    self.prob = prob\n",
    "    self.rng = bm.random.RandomState(seed=seed)\n",
    "\n",
    "  def update(self, x):\n",
    "    if bp.share.load('fit'):\n",
    "      keep_mask = self.rng.bernoulli(self.prob, x.shape)\n",
    "      return bm.where(keep_mask, x / self.prob, 0.)\n",
    "    else:\n",
    "      return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Here, the model makes different outputs according to the different values of a shared parameter ``fit``.\n",
    "\n",
    "You can define your own shared parameters, and then provide their shared parameters when calling the trainer objects (see the following section)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Examples of training models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "In the following, we illustrate several examples to build a trainable neural network model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Artificial neural networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "BrainPy provides neural network layers which can be useful to define artificial neural networks.\n",
    "\n",
    "Here, let's define a deep RNN model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.300765500Z",
     "start_time": "2023-04-15T17:18:08.625911Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class DeepRNN(bp.DynamicalSystemNS):\n",
    "    def __init__(self, num_in, num_recs, num_out):\n",
    "        super(DeepRNN, self).__init__()\n",
    "\n",
    "        self.l1 = bp.layers.LSTMCell(num_in, num_recs[0])\n",
    "        self.d1 = bp.layers.Dropout(0.2)\n",
    "        self.l2 = bp.layers.LSTMCell(num_recs[0], num_recs[1])\n",
    "        self.d2 = bp.layers.Dropout(0.2)\n",
    "        self.l3 = bp.layers.LSTMCell(num_recs[1], num_recs[2])\n",
    "        self.d3 = bp.layers.Dropout(0.2)\n",
    "        self.l4 = bp.layers.LSTMCell(num_recs[2], num_recs[3])\n",
    "        self.d4 = bp.layers.Dropout(0.2)\n",
    "        self.lout = bp.layers.Dense(num_recs[3], num_out)\n",
    "\n",
    "    def update(self, x):\n",
    "        x = x >> self.l2 >> self.d1\n",
    "        x = x >> self.l2 >> self.d2\n",
    "        x = x >> self.l3 >> self.d3\n",
    "        x = x >> self.l4 >> self.d4\n",
    "        return self.lout(x)\n",
    "\n",
    "with bm.training_environment():\n",
    "    model = DeepRNN(100, [200, 200, 200, 100], 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Note here the difference of the model building from PyTorch is that the first argument in ``update()`` function should be the shared parameters ``sha`` (i.e., these parameters are shared across all models, like the time ``t``, the running index ``i``, and the model running phase ``fit``). Then other individual arguments can all be customized by users. The details of the model definition specification can be seen in ????"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Moreover, it is worthy to note that this model only defines the one step updating rule of how the model evolves according to the input ``x``."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Reservoir computing models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "In this example, we define a reservoir computing model called [next generation reservoir computing](https://doi.org/10.1038/s41467-021-25801-2) by using the built-in models provided in BrainPy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.304527500Z",
     "start_time": "2023-04-15T17:18:09.308723Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class NGRC(bp.DynamicalSystemNS):\n",
    "  def __init__(self, num_in, num_out):\n",
    "    super(NGRC, self).__init__(mode=bm.batching_mode)\n",
    "    self.r = bp.layers.NVAR(num_in, delay=4, order=2, stride=5, mode=bm.batching_mode)\n",
    "    self.o = bp.layers.Dense(self.r.num_out, num_out, mode=bm.training_mode)\n",
    "\n",
    "  def update(self, x):\n",
    "    return x >> self.r >> self.o"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "In the above model, ``brainpy.layers.NVAR`` is a nonlinear vector autoregression machine, which does not have the training features. Therefore, we define its ``mode`` as batching mode. On the contrary, ``brainpy.layers.Dense`` has the trainable weights for model training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Spiking Neural Networks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Building trainable spiking neural networks in BrainPy is also a piece of cake. We provided commonly used spiking models for traditional dynamics simulation. But most of them can be used for training too.\n",
    "\n",
    "In the following, we provide an implementation of spiking neural networks in [(Neftci, Mostafa, & Zenke, 2019)](https://doi.org/10.1109/MSP.2019.2931595) for surrogate gradient learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-10-06T05:19:30.304527500Z",
     "start_time": "2023-04-15T17:18:09.327180Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class SNN(bp.Network):\n",
    "  def __init__(self, num_in, num_rec, num_out):\n",
    "    super(SNN, self).__init__()\n",
    "\n",
    "    # neuron groups\n",
    "    self.i = bp.neurons.InputGroup(num_in)\n",
    "    self.r = bp.neurons.LIF(num_rec, tau=10, V_reset=0, V_rest=0, V_th=1.)\n",
    "    self.o = bp.neurons.LeakyIntegrator(num_out, tau=5)\n",
    "\n",
    "    # synapse: i->r\n",
    "    self.i2r = bp.synapses.Exponential(self.i, self.r, bp.conn.All2All(),\n",
    "                                       output=bp.synouts.CUBA(), tau=10.,\n",
    "                                       g_max=bp.init.KaimingNormal(scale=20.))\n",
    "    # synapse: r->o\n",
    "    self.r2o = bp.synapses.Exponential(self.r, self.o, bp.conn.All2All(),\n",
    "                                       output=bp.synouts.CUBA(), tau=10.,\n",
    "                                       g_max=bp.init.KaimingNormal(scale=20.))\n",
    "\n",
    "  def update(self, tdi, spike):\n",
    "    self.i2r(tdi, spike)\n",
    "    self.r2o(tdi)\n",
    "    self.r(tdi)\n",
    "    self.o(tdi)\n",
    "    return self.o.V.value\n",
    "\n",
    "with bm.training_environment():\n",
    "    snn = SNN(10, 100, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Note here the mode in all models are specified as ``brainpy.modes.TrainingMode``."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
